# 导入包库
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable

import torchvision
import torchvision.transforms as transforms

from torch.utils.tensorboard import SummaryWriter


def get_num_correct(preds, labels):
    return preds.argmax(dim=1).eq(labels).sum().item()


class Network(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5)
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=12, kernel_size=5)
        self.fc1 = nn.Linear(in_features=12*4*4, out_features=120)
        self.fc2 = nn.Linear(in_features=120, out_features=60)
        self.out = nn.Linear(in_features=60, out_features=10)

    def forward(self, t):
        t = F.relu(self.conv1(t))
        t = F.max_pool2d(t, kernel_size=2, stride=2)

        t = F.relu(self.conv2(t))
        t = F.max_pool2d(t, kernel_size=2, stride=2)

        t = t.flatten(start_dim=1)
        t = F.relu(self.fc1(t))

        t = F.relu(self.fc2(t))
        t = self.out(t)

        return t


if __name__ == '__main__':
    # 我自己来设置路径
    from pathlib import Path
    PATH = Path(r'C:\files\git_repository\pytorch-learning\pytorch学习\通过示例学习')

    train_set = torchvision.datasets.FashionMNIST(
        root=str(PATH/'data'),
        train=True,
        download=True,
        transform=transforms.Compose([
            transforms.ToTensor()
        ])
    )

    train_loader = torch.utils.data.DataLoader(
        train_set, batch_size=100, shuffle=True)

    # tensor board
    tb = SummaryWriter()
    network = Network()
# 取出训练用图
    images, labels = next(iter(train_loader))
    grid = torchvision.utils.make_grid(images)
# 想用tensorboard看什么，你就tb.add什么。image、graph、scalar等
    tb.add_image('images', grid)
    tb.add_graph(model=network, input_to_model=images)
    tb.close()
    exit(0)
